import_path_ref = {
    "collective_assault": ("Environment.collective_assault.collective_assault_parallel_run", 'ScenarioConfig'),
    "dca_multiteam": ("Environment.dca_multiteam.collective_assault_parallel_run", 'ScenarioConfig'),
    "dca": ("Environment.dca.collective_assault_parallel_run", 'ScenarioConfig'),
    "air_fight": ("Environment.air_fight.environment.air_fight_compat", 'ScenarioConfig'),
    "native_gym": ("Environment.native_gym.native_gym_config", 'ScenarioConfig'),
    "starcraft2": ("Environment.starcraft.sc2_env_wrapper", 'ScenarioConfig'),
    "sc2": ("Environment.starcraft.sc2_env_wrapper", 'ScenarioConfig'),
    "unity_game": ("Environment.unity_game.unity_game_wrapper", 'ScenarioConfig'),
    "sr_tasks->cargo": ("Environment.sr_tasks.multiagent.scenarios.cargo", 'ScenarioConfig'),
    "sr_tasks->hunter_invader": ("Environment.sr_tasks.multiagent.scenarios.hunter_invader", 'ScenarioConfig'),
    "sr_tasks->hunter_invader3d": ("Environment.sr_tasks.multiagent.scenarios.hunter_invader3d", 'ScenarioConfig'),
    "sr_tasks->hunter_invader3d_v2": ("Environment.sr_tasks.multiagent.scenarios.hunter_invader3d_v2", 'ScenarioConfig'),
    "bvr": ("Environment.bvr_sim.init_env", 'ScenarioConfig'),
    "mathgame": ("Environment.math_game.env", 'ScenarioConfig'),
    # "cssim": ("Algorithm.Common.cssim_env_wrapper", 'ScenarioConfig'),

    # 2025/07：目录修改
    "cssim": ("Common.cssim_env_wrapper", 'ScenarioConfig'),
    "llm_trainer": ("Environment.llm_trainer.llm_env", 'ScenarioConfig'),
}

env_init_function_ref = {
    "collective_assault": ("Environment.collective_assault.collective_assault_parallel_run", 'make_collective_assault_env'),
    "dca_multiteam": ("Environment.dca_multiteam.collective_assault_parallel_run", 'make_collective_assault_env'),
    "dca": ("Environment.dca.collective_assault_parallel_run", 'make_collective_assault_env'),
    "air_fight": ("Environment.air_fight.environment.air_fight_compat", 'make_air_fight_env'),
    "native_gym": ("Environment.native_gym.native_gym_config", 'env_init_function'),
    "starcraft2": ("Environment.starcraft.sc2_env_wrapper", 'make_sc2_env'),
    "sc2": ("Environment.starcraft.sc2_env_wrapper", 'make_sc2_env'),
    "unity_game": ("Environment.unity_game.unity_game_wrapper", 'make_env'),
    "sr_tasks": ("Environment.sr_tasks.multiagent.scenario", 'sr_tasks_env'),
    "bvr": ("Environment.bvr_sim.init_env", 'make_bvr_env'),
    "mathgame": ("Environment.math_game.env", 'make_math_env'),
    # "cssim": ("Algorithm.Common.cssim_env_wrapper", 'make_uhmap_env'),

    # 2025/07:目录结构修改
    "cssim": ("Common.cssim_env_wrapper", 'make_uhmap_env'),
    "uhmap_game": ("Algorithm.Common.cssim_env_wrapper", 'make_game_env'),
    "llm_trainer": ("Environment.llm_trainer.llm_env", 'make_llm_env'),
}

##################################################################################################################################
##################################################################################################################################
from config import GlobalConfig
import importlib, os
# from Utils.colorful import print亮蓝


def load_ScenarioConfig():
    if GlobalConfig.env_name not in import_path_ref:
        assert False, ('need to find path of ScenarioConfig')
    import_path, ScenarioConfig = import_path_ref[GlobalConfig.env_name]
    GlobalConfig.ScenarioConfig = getattr(importlib.import_module(import_path), ScenarioConfig)


def load_AlgorithmConfig():     # 2024/08 修改

    # 2025/07
    assert hasattr(GlobalConfig, 'AlgorithmConfig_Red'), (
        GlobalConfig, 'has no such config item: **%s**' % 'AlgorithmConfig_Red')

    assert hasattr(GlobalConfig, 'AlgorithmConfig_Blue'), (
        GlobalConfig, 'has no such config item: **%s**' % 'AlgorithmConfig_Blue')

    # load Red AlgorithmConfig 2025/07
    algorithmConfig_path = GlobalConfig.AlgorithmConfig_Red
    if isinstance(algorithmConfig_path, str) and '->' in algorithmConfig_path:  # 2024/08 修改
        # import importlib
        file_, class_ = algorithmConfig_path.split('->')
        if '.py' in file_:
            if file_.endswith('.py'): file_ = file_[:-3]
        GlobalConfig.AlgorithmConfig_Red = getattr(importlib.import_module(file_), class_)

    # load Blue AlgorithmConfig 2025/07
    algorithmConfig_path = GlobalConfig.AlgorithmConfig_Blue
    if isinstance(algorithmConfig_path, str) and '->' in algorithmConfig_path:  # 2024/08 修改
        # import importlib
        file_, class_ = algorithmConfig_path.split('->')
        if '.py' in file_:
            if file_.endswith('.py'): file_ = file_[:-3]
        GlobalConfig.AlgorithmConfig_Blue = getattr(importlib.import_module(file_), class_)

def make_env_function(env_name, rank, algorithms):
    load_ScenarioConfig()
    ref_env_name = env_name

    if 'native_gym' in env_name:
        assert '->' in env_name
        ref_env_name, env_name = env_name.split('->')
    elif 'sr_tasks' in env_name:
        assert '->' in env_name
        ref_env_name, env_name = env_name.split('->')

    import_path, func_name = env_init_function_ref[ref_env_name]
    env_init_function = getattr(importlib.import_module(import_path), func_name)
    return lambda: env_init_function(env_name, rank, algorithms)

def load_algorithm_list():
    n_t = GlobalConfig.ScenarioConfig.N_TEAM  # n_t => n_teams
    n_agents_each_t = GlobalConfig.ScenarioConfig.N_AGENT_EACH_TEAM  # n_agents_each_t => n_agents_each_team
    t_member_list = GlobalConfig.ScenarioConfig.AGENT_ID_EACH_TEAM
    t_name = GlobalConfig.ScenarioConfig.Algorithm_Names

    algo_foundations = []  # import and initialize algorithms
    for t in range(n_t):
        assert len(t_member_list[t]) == n_agents_each_t[t]
        assert '->' in t_name[t]
        module_, class_ = t_name[t].split('->')
        init_f = getattr(importlib.import_module(module_), class_)
        algo_foundations.append(init_f)

    return algo_foundations

def make_parallel_envs(process_pool, algorithms, marker='',):
    from Utils.shm_env import SuperpoolEnv
    from config import GlobalConfig
    # from Environment.Common.route_and_environment import load_ScenarioConfig

    # 2025/07: Move to task_runner.py ----modify begin
    # load_ScenarioConfig()
    # load_AlgorithmConfig()  # 2024/08 修改
    # 2025/07: Move to task_runner.py ----modify end

    env_args_dict_list = [({
                               'env_name': GlobalConfig.env_name,
                               'proc_index': i if 'test' not in marker else -(i + 1),
                               'marker': marker,
                               'algorithms':algorithms  # 2025/07:In order to get algorithm's reward function in envsS
                           },) for i in range(GlobalConfig.num_worker_threads)]

    # if GlobalConfig.env_name == 'air_fight':
    #     # This particular env has a dll file
    #     # that must be loaded in main process
    #     from Algorithm.air_fight.environment.pytransform import pyarmor_runtime
    #     pyarmor_runtime()

    if GlobalConfig.env_name == 'bvr':
        # 1、如果没用hmp的docker，请设置好 YOUR_ROOT_PASSWORD，不止这一处，请全局搜索"YOUR_ROOT_PASSWORD"替换所有
        # 2、用docker的sock挂载到容器中，方法在SetupDocker.md中
        print('[env_router]: here goes the docker in docker check.')
        YOUR_ROOT_PASSWORD = 'clara'  # the sudo password
        os.system("echo %s|sudo -S date" % YOUR_ROOT_PASSWORD)  # get sudo power
        res = os.popen("sudo docker ps").read()
        if "CONTAINER ID" not in res:
            print('[env_router]: Error checking docker in docker, can not control host docker interface!')
            raise "Error checking docker in docker, can not control host docker interface!"
        pass

    # if GlobalConfig.env_name == 'dca':
    #     # This particular env has a cython file that needs to be compiled in main process
    #     # that must be loaded in main process
    #     from Environment.dca.cython_func import laser_hit_improve3
    # if GlobalConfig.env_name == 'dca_multiteam':
    #     # This particular env has a cython file that needs to be compiled in main process
    #     # that must be loaded in main process
    #     from Environment.dca_multiteam.cython_func import laser_hit_improve3
    if GlobalConfig.env_name == 'cssim':
        # This particular env has a cython file that needs to be compiled in main process
        # that must be loaded in main process
        from Common.cython_env_func import tear_number_apart

    if GlobalConfig.num_worker_threads > 1:
        envs = SuperpoolEnv(process_pool, env_args_dict_list)
    else:
        envs = SuperpoolEnv(process_pool, env_args_dict_list)

    return envs
